import pandas as pd
import requests
import json
def load_data(file_path):
    # return pd.read_excel(file_path)
    return pd.read_csv(file_path)

def preprocess_data(df):
    df1 = df.dropna(subset=['问题'])
    df2 = df1.dropna(subset=['正确答案'])
    return df2

def analyze_questions(df):
    Q, R = [], []
    for i in range(len(df)):
        try:
            question = df['问题'].iloc[i]
            right_answer = df['正确答案'].iloc[i]
            Q.append(question)
            R.append(right_answer)
        except Exception as e:
            print(f'Error at index {i+1}: {e}')
            continue
    return Q, R

def query_responses(questions):
    url = 'http://172.31.30.3:8777/api/local_doc_qa/local_doc_chat'
    A, SOURCE = [], []
    for i, question in enumerate(questions):
        headers = {"content-type": "application/json"}
        data = {
            "user_id": "zzp",
            "kb_ids": ["KB163dd069627044bca421e1ecfa856260"],
            "question": question,
            "history": [],
            "rerank": "true",
            # "streaming": "false"
        }
        response = requests.request("POST", url, json=data, headers=headers).json()
        # response = requests.post(url, json=data,headers=headers).json()
        if response["code"] == 200:        
            ai_response = response['response']
            content = [i["content"] for i in response["source_documents"]]
            # source_info = '\n'.join([doc['kernel'].split('KB_CONTENT: ')[-1] for doc in response.json().get('source_documents', [])])
            SOURCE.append(content)
            print(f'{i+1}条:', ai_response)
            A.append(ai_response)
        else:
            print(('请求失败，状态码:', response.status_code))
    return A, SOURCE

def save_result(df, file_path):
    df.to_csv(file_path, index=False, encoding='utf8')

def main():
    # test_path = r'/data/lrs/QAnything/qanything_kernel/test_data/test0305new.xlsx'
    test_path = r'/workspace/qanything_local/qanything_kernel/test_data/test0305new.csv'
    test_result_path = r'/workspace/qanything_local/qanything_kernel/test_data/1test_result_0312.csv'

    df = load_data(test_path)
    
    df_processed = preprocess_data(df)

    Q, R = analyze_questions(df_processed)

    A, SOURCE = query_responses(Q)

    df['LLM生成的答案'] = A
    df['溯源信息'] = SOURCE

    save_result(df, test_result_path)

if __name__ == "__main__":
    main()
